
# Figure 1

library(dplyr)
library(ggplot2)



#These functions find the frontier values of t and p, meaning the values that give us the upper envelope of the HL-HW space
t_frontier<-function(alpha=1.25){
  alpha/(alpha-1) - 1
}

p_frontier<-function(K,I=10,alpha=1.25,w=10){
  val <- w-((-K+I+log(t_frontier(alpha)+1))/(alpha-1))
  
  ifelse(val>0,sqrt(val),0)
}


#I=10, alpha=1.25, ell(t)=log(t+1), w=10, A=w-p^2,

#Target point has p^*=1.25
#alpha=1.25

A_fun<-function(p,w=10){w-(p)^2}
HL<-function(p,t,alpha=1.25,I=10,w=10){I+(1-alpha)*A_fun(p,w) + log(t+1)}
HW<-function(p,t,alpha=1.25,I=10,w=10){I+alpha*A_fun(p,w) - t}

v_t<-c(log(t_frontier()+1),-t_frontier())
v_p<-(A_fun(p=1.25)-A_fun(0))*c(1-1.25,1.25)

displacement<-setNames(v_t+v_p,c("hl","hw"))

norm_euclid<-function(vec_end,vec_start){
  sqrt(sum((vec_end-vec_start)^2))
}

df2<-expand.grid(p=seq(0,sqrt(10),0.01),t=seq(0,5,0.01)) %>% 
  mutate(hw=HW(p,t),hl=HL(p,t)) %>% 
  filter(hl<=10) %>%
  mutate(hl_tile=ntile(hl,50)) %>%
  group_by(hl_tile) %>%
  mutate(
    hl_mean=mean(hl),
    hw_min=min(hw),
    hw_max=max(hw)) %>%
  ungroup() %>%
  arrange(hl) %>%
  mutate(hl_num=1:n())


origin<-setNames(as.numeric(df2[df2$p==0 & df2$t==0,c("hl","hw")]),c("hl","hw"))
endpoint<-origin+displacement


plt.cosmopolitan<-ggplot()+
  scale_color_viridis_c(option="cividis")+
  theme_bw()+
  theme(legend.position="bottom")+
  labs(color="t")+
  #v_t
  geom_segment(
    data=data.frame(
      x=c(origin["hl"],origin["hl"]+(displacement["hl"]-log(t_frontier()+1))),
      y=c(origin["hw"],origin["hw"]-((origin["hw"]-t_frontier())-endpoint["hw"])),
      xend=c(origin["hl"]+log(t_frontier()+1),origin["hl"]+log(t_frontier()+1)+(displacement["hl"]-log(t_frontier()+1))),
      yend=c(origin["hw"]-t_frontier(),origin["hw"]-t_frontier()-((origin["hw"]-t_frontier())-endpoint["hw"]))
    )[1,],
    aes(
      x=x,
      y=y,
      xend=xend,
      yend=yend),
    color=khroma::color(palette = "bright")(4)[1],
    arrow=arrow(type="open",angle=15),
    arrow.fill="black",
    size=1)+
  #v_p
  geom_segment(
    data=data.frame(
      x=c(origin["hl"]+log(t_frontier()+1),origin["hl"]),
      y=c(origin["hw"]-t_frontier(),origin["hw"]),
      xend=c(endpoint["hl"],origin["hl"]+(displacement["hl"]-log(t_frontier()+1))),
      yend=c(endpoint["hw"],origin["hw"]-((origin["hw"]-t_frontier())-endpoint["hw"]))
    )[1,],
    aes(
      x=x,
      y=y,
      xend=xend,
      yend=yend),
    color=khroma::color(palette = "bright")(4)[2],
    arrow=arrow(type="open",angle=15),
    arrow.fill="black",
    size=1)+
  #v_t+v_p
  geom_segment(
    aes(
      x=origin["hl"],
      y=origin["hw"],
      xend=endpoint["hl"],
      yend=endpoint["hw"]),
    color=khroma::color(palette = "bright")(4)[3],
    arrow=arrow(type="open",angle=15),
    arrow.fill="black",
    size=1)+
  geom_label(
    mapping=
      aes(
        x=(origin["hl"]+endpoint["hl"])/2,
        y=(origin["hw"]+endpoint["hw"])/2
      ),
    label=eval(substitute(expression("||"*v["t"]+v["p"]*"||" == A),list(A=round(norm_euclid(endpoint,origin),2)))),
    color=khroma::color(palette = "bright")(4)[3],
    nudge_x = -0.25,
    nudge_y = -2
  )+
  geom_segment(
    aes(
      x=origin["hl"],
      y=origin["hw"],
      xend=(v_t %*% displacement)/(displacement %*% displacement) * displacement["hl"] + origin["hl"],
      yend=(v_t %*% displacement)/(displacement %*% displacement) * displacement["hw"] + origin["hw"]
    ),
    color=khroma::color(palette = "bright")(4)[4],
    arrow=arrow(angle=15,type="open")
  )+
  geom_label(
    mapping=
      aes(
        x=(origin["hl"]+(v_t %*% displacement)/(displacement %*% displacement) * displacement["hl"] + origin["hl"])/2,
        y=(origin["hw"]+(v_t %*% displacement)/(displacement %*% displacement) * displacement["hw"] + origin["hw"])/2
      ),
    label=eval(substitute(expression("||"*italic(proj)[v["t"]+v["p"]](v["t"])*"||" == A),list(A=round(norm_euclid(c((v_t %*% displacement)/(displacement %*% displacement) * displacement["hl"] + origin["hl"],(v_t %*% displacement)/(displacement %*% displacement) * displacement["hw"] + origin["hw"]),origin),2)))),
    color=khroma::color(palette = "bright")(4)[4],
    nudge_x=0.5,
    nudge_y=0.75
  )+
  labs(x=expression(paste("Income of Group L (",H["L"],")")),y=expression(paste("Income of Group W (",H["W"],")")),title="Cosmopolitan \n Any Shock")


####### NATIONALISTS


#imagine that the nationalist increases their p by 20% and decreases their t by 15% 

t_n<-0.25*t_frontier()

v_t_n<-setNames(c(log(t_n+1),-t_n),c("hl","hw"))
v_p_n<-setNames((A_fun(p=1.20*1.25)-A_fun(0))*c(1-(1.25),1.25),c("hl","hw"))

displacement_n<-setNames(v_t_n+v_p_n,c("hl","hw"))

origin<-setNames(as.numeric(df2[df2$p==0 & df2$t==0,c("hl","hw")]),c("hl","hw"))
endpoint_n<-origin+displacement_n


plt.nationalist<-ggplot()+
  scale_color_viridis_c(option="cividis")+
  theme_bw()+
  theme(legend.position="bottom")+
  labs(color="t")+
  #v_t
  geom_segment(
    data=data.frame(
      x=c(origin["hl"],origin["hl"]+v_p_n["hl"]),
      y=c(origin["hw"],origin["hw"]+v_p_n["hw"]),
      xend=c(origin["hl"]+v_t_n["hl"],origin["hl"]+v_p_n["hl"]+v_t_n["hl"]),
      yend=c(origin["hw"]+v_t_n["hw"],origin["hw"]+v_p_n["hw"]+v_t_n["hw"])
    )[1,],
    aes(
      x=x,
      y=y,
      xend=xend,
      yend=yend),
    color=khroma::color(palette = "bright")(4)[1],
    arrow=arrow(type="open",angle=15),
    arrow.fill="black",
    size=1)+
  #v_p
  geom_segment(
    data=data.frame(
      x=c(origin["hl"]+v_t_n["hl"],origin["hl"]),
      y=c(origin["hw"]+v_t_n["hw"],origin["hw"]),
      xend=c(endpoint_n["hl"],origin["hl"]+v_p_n["hl"]),
      yend=c(endpoint_n["hw"],origin["hw"]+v_p_n["hw"])
    )[1,],
    aes(
      x=x,
      y=y,
      xend=xend,
      yend=yend),
    color=khroma::color(palette = "bright")(4)[2],
    arrow=arrow(type="open",angle=15),
    arrow.fill="black",
    size=1)+
  #v_t+v_p
  geom_segment(
    aes(
      x=origin["hl"],
      y=origin["hw"],
      xend=endpoint_n["hl"],
      yend=endpoint_n["hw"]),
    color=khroma::color(palette = "bright")(4)[3],
    arrow=arrow(type="open",angle=15),
    arrow.fill="black",
    size=1)+
  geom_label(
    mapping=
      aes(
        x=(origin["hl"]+endpoint_n["hl"])/2,
        y=(origin["hw"]+endpoint_n["hw"])/2
      ),
    label=eval(substitute(expression("||"*v["t"]+v["p"]*"||" == A),list(A=round(norm_euclid(endpoint_n,origin),2)))),
    color=khroma::color(palette = "bright")(4)[3],
    nudge_x = 0,
    nudge_y = -2
  )+
  geom_segment(
    aes(
      x=origin["hl"],
      y=origin["hw"],
      xend=(v_t_n %*% displacement_n)/(displacement_n %*% displacement_n) * displacement_n["hl"] + origin["hl"],
      yend=(v_t_n %*% displacement_n)/(displacement_n %*% displacement_n) * displacement_n["hw"] + origin["hw"]
    ),
    color=khroma::color(palette = "bright")(4)[4],
    arrow=arrow(angle=15,type="open")
  )+
  geom_label(
    mapping=
      aes(
        x=(origin["hl"]+(v_t_n %*% displacement_n)/(displacement_n %*% displacement_n) * displacement_n["hl"] + origin["hl"])/2,
        y=(origin["hw"]+(v_t_n %*% displacement_n)/(displacement_n %*% displacement_n) * displacement_n["hw"] + origin["hw"])/2
      ),
    label=eval(substitute(expression("||"*italic(proj)[v["t"]+v["p"]](v["t"])*"||" == A),list(A=round(norm_euclid(c((v_t_n %*% displacement_n)/(displacement_n %*% displacement_n) * displacement_n["hl"] + origin["hl"],(v_t_n %*% displacement_n)/(displacement_n %*% displacement_n) * displacement_n["hw"] + origin["hw"]),origin),2)))),
    color=khroma::color(palette = "bright")(4)[4],
    nudge_x=0.75,
    nudge_y=0.5
  )+
  labs(x=expression(paste("Income of Group L (",H["L"],")")),y=expression(paste("Income of Group W (",H["W"],")")),title="Nationalist \n Foreign Shock")


##### NATIONALIST 2 

#imagine that the nationalist decreases their p by 20% and increases their t by 15% 

t_n_2<-1.05*t_frontier()

v_t_n_2<-setNames(c(log(t_n_2+1),-t_n_2),c("hl","hw"))
v_p_n_2<-setNames((A_fun(p=0.8*1.25)-A_fun(0))*c(1-(1.25),1.25),c("hl","hw"))

displacement_n_2<-setNames(v_t_n_2+v_p_n_2,c("hl","hw"))

origin<-setNames(as.numeric(df2[df2$p==0 & df2$t==0,c("hl","hw")]),c("hl","hw"))
endpoint_n_2<-origin+displacement_n_2


plt.nationalist.2<-ggplot()+
  scale_color_viridis_c(option="cividis")+
  theme_bw()+
  theme(legend.position="bottom")+
  labs(color="t")+
  #v_t
  geom_segment(
    data=data.frame(
      x=c(origin["hl"],origin["hl"]+v_p_n_2["hl"]),
      y=c(origin["hw"],origin["hw"]+v_p_n_2["hw"]),
      xend=c(origin["hl"]+v_t_n_2["hl"],origin["hl"]+v_p_n_2["hl"]+v_t_n_2["hl"]),
      yend=c(origin["hw"]+v_t_n_2["hw"],origin["hw"]+v_p_n_2["hw"]+v_t_n_2["hw"])
    )[1,],
    aes(
      x=x,
      y=y,
      xend=xend,
      yend=yend),
    color=khroma::color(palette = "bright")(4)[1],
    arrow=arrow(type="open",angle=15),
    arrow.fill="black",
    size=1)+
  #v_p
  geom_segment(
    data=data.frame(
      x=c(origin["hl"]+v_t_n_2["hl"],origin["hl"]),
      y=c(origin["hw"]+v_t_n_2["hw"],origin["hw"]),
      xend=c(endpoint_n_2["hl"],origin["hl"]+v_p_n_2["hl"]),
      yend=c(endpoint_n_2["hw"],origin["hw"]+v_p_n_2["hw"])
    )[1,],
    aes(
      x=x,
      y=y,
      xend=xend,
      yend=yend),
    color=khroma::color(palette = "bright")(4)[2],
    arrow=arrow(type="open",angle=15),
    arrow.fill="black",
    size=1)+
  #v_t+v_p
  geom_segment(
    aes(
      x=origin["hl"],
      y=origin["hw"],
      xend=endpoint_n_2["hl"],
      yend=endpoint_n_2["hw"]),
    color=khroma::color(palette = "bright")(4)[3],
    arrow=arrow(type="open",angle=15),
    arrow.fill="black",
    size=1)+
  geom_label(
    mapping=
      aes(
        x=(origin["hl"]+endpoint_n_2["hl"])/2,
        y=(origin["hw"]+endpoint_n_2["hw"])/2
      ),
    label=eval(substitute(expression("||"*v["t"]+v["p"]*"||" == A),list(A=round(norm_euclid(endpoint_n_2,origin),2)))),
    color=khroma::color(palette = "bright")(4)[3],
    nudge_x = 0,
    nudge_y = -2
  )+
  geom_segment(
    aes(
      x=origin["hl"],
      y=origin["hw"],
      xend=(v_t_n_2 %*% displacement_n_2)/(displacement_n_2 %*% displacement_n_2) * displacement_n_2["hl"] + origin["hl"],
      yend=(v_t_n_2 %*% displacement_n_2)/(displacement_n_2 %*% displacement_n_2) * displacement_n_2["hw"] + origin["hw"]
    ),
    color=khroma::color(palette = "bright")(4)[4],
    arrow=arrow(angle=15,type="open")
  )+
  geom_label(
    mapping=
      aes(
        x=(origin["hl"]+(v_t_n_2 %*% displacement_n_2)/(displacement_n_2 %*% displacement_n_2) * displacement_n_2["hl"] + origin["hl"])/2,
        y=(origin["hw"]+(v_t_n_2 %*% displacement_n_2)/(displacement_n_2 %*% displacement_n_2) * displacement_n_2["hw"] + origin["hw"])/2
      ),
    label=eval(substitute(expression("||"*italic(proj)[v["t"]+v["p"]](v["t"])*"||" == A),list(A=round(norm_euclid(c((v_t_n_2 %*% displacement_n_2)/(displacement_n_2 %*% displacement_n_2) * displacement_n_2["hl"] + origin["hl"],(v_t_n_2 %*% displacement_n_2)/(displacement_n_2 %*% displacement_n_2) * displacement_n_2["hw"] + origin["hw"]),origin),2)))),
    color=khroma::color(palette = "bright")(4)[4],
    nudge_x=0.4,
    nudge_y=1.25
  )+
  labs(x=expression(paste("Income of Group L (",H["L"],")")),y=expression(paste("Income of Group W (",H["W"],")")),title="Nationalist \n Domestic Shock")


#Output

gridExtra::grid.arrange(
  plt.cosmopolitan+coord_equal(xlim=c(7.5,9.5),ylim=c(16.5,22.5)),
  plt.nationalist+coord_equal(xlim=c(7.5,9.5),ylim=c(16.5,22.5)),
  plt.nationalist.2+coord_equal(xlim=c(7.5,9.5),ylim=c(16.5,22.5)),
  ncol=3)

